{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Learn OpenAI Whisper - Chapter 5\n",
        "## Notebook 3: Voice assistant with Whisper and StableLM Zephyr\n",
        "\n",
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1LPRZl1JfLfhDFP75XkCMFJbDqQMWvY2L)\n",
        "\n",
        "This notebook was designed for setting up and demonstrating a voice assistant system that integrates OpenAI's Whisper model for speech-to-text capabilities and the [StableLM Zephyr 3B](https://huggingface.co/TheBloke/stablelm-zephyr-3b-GGUF) model for natural language understanding and response generation. The notebook is structured to facilitate the installation of necessary libraries, loading of the models, and execution of the voice assistant functionality, leveraging GPU acceleration for improved performance.\n",
        "\n",
        "### Main code sections and their functions\n",
        "\n",
        "The notebook is structured into several main code sections, each with a specific function:\n",
        "\n",
        "1. **Setting up the environment**: This section sets environmental variables and installs necessary Python packages such as `llama-cpp-python`, `whisper`, `gradio`, and `gTTS` for text-to-speech conversion.\n",
        "\n",
        "2. **Initializating Python libraries**: It imports essential libraries and sets up a logger to record events and outputs during the notebook's execution.\n",
        "\n",
        "3. **Loading the inference model**: This part of the notebook loads the StableLM Zephyr 3B model with LLAMA.CPP, configuring it for GPU usage if available. It specifies parameters such as the maximum sequence length, number of CPU threads, and the number of layers to offload to the GPU.\n",
        "\n",
        "4. **Exploring an inference example**: A simple example is provided to demonstrate how to generate a response from the StableLM Zephyr 3B model given a text prompt.\n",
        "\n",
        "5. **Defining supporting functions**: Functions are defined for interacting with the StableLM model (`llm_call`) and for transcribing audio input into text using Whisper, then generating a response and converting it to speech (`transcribe`).\n",
        "\n",
        "6. **Creating an audio 'seed'**: A temporary audio file is created using `ffmpeg` to facilitate audio handling within the notebook.\n",
        "\n",
        "7. **Running a Gradio interface**: The final section sets up a Gradio interface, which provides a user-friendly GUI for real-time interaction with the voice assistant. It defines the inputs and outputs for the interface and launches it for use.\n",
        "\n",
        "The programming language used in the Jupyter notebook is Python, as indicated by the kernel specification and the language_info metadata in the notebook file."
      ],
      "metadata": {
        "id": "wK84o9zKU5uo"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "1. **Setting up the environment**\n",
        "    \n",
        "    This section sets environmental variables and installs necessary Python packages such as `llama-cpp-python`, `whisper`, `gradio`, and `gTTS` for text-to-speech conversion.\n",
        "\n",
        "    **Setting up environmental variable for llama.ccp**\n",
        "\n",
        "    Before loading the `stablelm-zephyr-3b-GGUF stablelm-zephyr-3b.Q5_K_S.gguf` model, we need to install and compile the `llama-cpp-python` package. To leverage NVIDIA CUDA acceleration, we must set an environmental variable\n",
        "`CMAKE_ARGS=\"-DLLAMA_CUBLAS=on\"` first."
      ],
      "metadata": {
        "id": "gWjGJgc2XtNU"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sOi3PQHTwtrk"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "os.environ[\"CMAKE_ARGS\"] = \"-DLLAMA_CUBLAS=on\"\n",
        "print(os.getenv(\"CMAKE_ARGS\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6bvyny_mf2na"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install llama-cpp-python==0.2.34\n",
        "!huggingface-cli download TheBloke/stablelm-zephyr-3b-GGUF stablelm-zephyr-3b.Q5_K_S.gguf --local-dir . --local-dir-use-symlinks False\n",
        "!pip install -q git+https://github.com/openai/whisper.git\n",
        "!pip install -q gradio\n",
        "!pip install -q gTTS"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "2. **Initializating Python libraries**\n",
        "\n",
        "    We now import essential libraries and set up a logger to record events and outputs during the notebook's execution."
      ],
      "metadata": {
        "id": "_VKYJTBtZ6Uj"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I859xTaaiCZ_"
      },
      "outputs": [],
      "source": [
        "import datetime\n",
        "import os\n",
        "from rich.console import Console\n",
        "console = Console(width=110)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7DcwSW-8izb3"
      },
      "outputs": [],
      "source": [
        "## Logger file\n",
        "tstamp = datetime.datetime.now()\n",
        "tstamp = str(tstamp).replace(' ','_')\n",
        "logfile = f'{tstamp}_log.txt'\n",
        "def writehistory(text):\n",
        "    with open(logfile, 'a', encoding='utf-8') as f:\n",
        "        f.write(text)\n",
        "        f.write('\\n')\n",
        "    f.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "3. **Loading the inference model**\n",
        "\n",
        "    This section of the notebook loads the StableLM Zephyr 3B model with LLAMA.CPP, configuring it for GPU usage if available. It specifies parameters such as the maximum sequence length, number of CPU threads, and the number of layers to offload to the GPU."
      ],
      "metadata": {
        "id": "ZnpRS8D7aDtu"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gp0dQJCKqkSr"
      },
      "outputs": [],
      "source": [
        "## Load a llama-cpp-python quantized model\n",
        "from llama_cpp import Llama\n",
        "with console.status(\"Loading ✅✅✅✅ stablelm-zephyr-3b with LLAMA.CPP...\",spinner=\"dots12\"):\n",
        "  llm_gpu = Llama(\n",
        "    model_path=\"/content/stablelm-zephyr-3b.Q5_K_S.gguf\",  # Download the model file first\n",
        "    n_ctx=4096,  # The max sequence length to use - note that longer sequence lengths require much more resources\n",
        "    n_threads=8,            # The number of CPU threads to use, tailor to your system and the resulting performance\n",
        "    n_gpu_layers=35         # The number of layers to offload to GPU, if you have GPU acceleration available\n",
        ")\n",
        "writehistory(f\"{str(datetime.datetime.now())} Loaded 🧠 stablelm-zephyr-3b.Q5_K_S.gguf with GPU enabled\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "4. **Exploring an inference example**\n",
        "    \n",
        "    A simple example is provided to demonstrate how to generate a response from the StableLM Zephyr 3B model given a text prompt."
      ],
      "metadata": {
        "id": "oX2jLkxAaOtR"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IdmxO4GMq28R"
      },
      "outputs": [],
      "source": [
        "# Simple inference example\n",
        "prompt=\"In short response, what is the capital of France?\"\n",
        "\n",
        "template = f\"<|user|>\\n{prompt}<|endoftext|>\\n<|assistant|>\"\n",
        "\n",
        "start = datetime.datetime.now()\n",
        "output = llm_gpu(\n",
        "    template, # Prompt\n",
        "    temperature=0,\n",
        "    max_tokens=512,  # Generate up to 512 tokens\n",
        "    stop=[\"</s>\"],   # Example stop token - not necessarily correct for this specific model! Please check before using.\n",
        "    echo=False        # Whether to echo the prompt\n",
        ")\n",
        "delta = datetime.datetime.now() - start\n",
        "console.print(f\"[bright_green bold on black]Question: {prompt}\")\n",
        "console.print(output['choices'][0]['text'])\n",
        "console.print(f\"Completed in: [bold red]{delta}\")\n",
        "writehistory(f\"{str(datetime.datetime.now())} Inference completed in: {delta}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "5. **Defining a supporting functions for the LLM**\n",
        "\n",
        "    Here we create a function for interacting with the StableLM model and test it."
      ],
      "metadata": {
        "id": "f69McytTaWUh"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uSNRY2N0H5_7"
      },
      "outputs": [],
      "source": [
        "import re\n",
        "\n",
        "def llm_call(input_text):\n",
        "    prompt = \"\"\"Act as Tatianna, a junior-level assistant characterized by your cheerful demeanor and unwavering helpfulness. \\\n",
        "    You are in a business setting, thus always act professionally and courteously. \\\n",
        "    Respond succinctly to the following instructions and questions, and do not include information about yourself unless it is part of the action or question: \\\n",
        "    \"\"\" + input_text\n",
        "\n",
        "    template = f\"<|user|>\\n{prompt}<|endoftext|>\\n<|assistant|>\"\n",
        "\n",
        "    start = datetime.datetime.now()\n",
        "    response = llm_gpu(\n",
        "        template, # Prompt\n",
        "        temperature=0.1,\n",
        "        max_tokens=200,  # Generate up to 512 tokens\n",
        "        stop=[\"</s>\"],   # Example stop token - not necessarily correct for this specific model! Please check before using.\n",
        "        echo=False        # Whether to echo the prompt\n",
        "    )\n",
        "\n",
        "    delta = datetime.datetime.now() - start\n",
        "    writehistory(f\"{str(datetime.datetime.now())} Inference completed in: {delta}\")\n",
        "\n",
        "    if response is not None:\n",
        "        match = re.search(r':\\s*(.*)', response['choices'][0]['text'])\n",
        "        if match:\n",
        "            reply = match.group(1).strip()\n",
        "        reply = response['choices'][0]['text']\n",
        "    else:\n",
        "        reply = \"No response generated.\"\n",
        "    return reply"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qR5ESVFzSPAM"
      },
      "outputs": [],
      "source": [
        "llm_call(\"Hello, what is your name, and can you bring me some coffee?\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "6. **Loading Whisper and creating a function for transcription**\n"
      ],
      "metadata": {
        "id": "RM-UeiKWbuXf"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vXULc7EKVFLC"
      },
      "outputs": [],
      "source": [
        "import warnings\n",
        "from gtts import gTTS\n",
        "import numpy as np\n",
        "import torch\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "torch.cuda.is_available()\n",
        "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "print(f\"Using torch {torch.__version__} ({DEVICE})\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MdTWHxPzVP4W"
      },
      "outputs": [],
      "source": [
        "import whisper\n",
        "model = whisper.load_model(\"medium\", device=DEVICE)\n",
        "print(\n",
        "    f\"Model is {'multilingual' if model.is_multilingual else 'English-only'} \"\n",
        "    f\"and has {sum(np.prod(p.shape) for p in model.parameters()):,} parameters.\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Bc3pDc1jUpv5"
      },
      "outputs": [],
      "source": [
        "def transcribe(audio):\n",
        "    # Check if the audio input is None or empty\n",
        "    if audio is None or audio == '':\n",
        "        return ('','',None)  # Return empty strings and None audio file\n",
        "\n",
        "    language = 'en'\n",
        "\n",
        "    audio = whisper.load_audio(audio)\n",
        "    audio = whisper.pad_or_trim(audio)\n",
        "\n",
        "    mel = whisper.log_mel_spectrogram(audio).to(model.device)\n",
        "\n",
        "    _, probs = model.detect_language(mel)\n",
        "\n",
        "    options = whisper.DecodingOptions()\n",
        "    result = whisper.decode(model, mel, options)\n",
        "    result_text = result.text\n",
        "\n",
        "    out_result = llm_call(result_text)\n",
        "\n",
        "    audioobj = gTTS(text = out_result,\n",
        "                    lang = language,\n",
        "                    slow = False)\n",
        "\n",
        "    audioobj.save(\"Temp.mp3\")\n",
        "\n",
        "    return [result_text, out_result, \"Temp.mp3\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "7. **Creating an audio 'seed'**\n",
        "\n",
        "    A temporary audio file is created using `ffmpeg` to facilitate audio handling within the notebook.\n"
      ],
      "metadata": {
        "id": "BpLxmWsSa6eE"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rO3yejFwHzGq"
      },
      "outputs": [],
      "source": [
        "# !ffmpeg -f lavfi -i anullsrc=r=44100:cl=mono -t 10 -q:a 9 -acodec libmp3lame Temp.mp3"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "8. **Running a Gradio interface**\n",
        "\n",
        "    The final section sets up a Gradio interface, which provides a user-friendly GUI for real-time interaction with the voice assistant. It defines the inputs and outputs for the interface and launches it for use."
      ],
      "metadata": {
        "id": "ALszLdH6cQnl"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CEp-FeYCV0Yk"
      },
      "outputs": [],
      "source": [
        "import gradio as gr\n",
        "\n",
        "output_1 = gr.Textbox(label=\"Speech to Text\")\n",
        "output_2 = gr.Textbox(label=\"ChatGPT Output\")\n",
        "output_3 = gr.Audio(\"Temp.mp3\", autoplay=True)\n",
        "\n",
        "gr.Interface(\n",
        "    title = 'Learn OpenAI Whisper: Voice Assistant - Using the StableLM Zephyr 3B model',\n",
        "    fn=transcribe,\n",
        "    # gr.inputs.Audio(source=\"microphone\", type=\"filepath\")\n",
        "    inputs = gr.Audio(sources=[\"microphone\"], type=\"filepath\"),\n",
        "    outputs=[\n",
        "        output_1,  output_2, output_3\n",
        "    ],\n",
        "    live=True).launch(debug=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Gratitude\n",
        "\n",
        "Many thanks to Naval Katoch for his valuable technical insights."
      ],
      "metadata": {
        "id": "UFnw_fiCxkMR"
      }
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}